#!/usr/bin/env python3
# encoding: utf-8
# coding style: pep8
# ====================================================
#   Copyright (C) 2021  All rights reserved.
#
#   Author        : Jinjing Li
#   Email         : josiahleeoaa@outlook.com
#   File Name     : plot_alphan.py
#   Last Modified : 2021-11-26 00:50
#   Describe      :
#
# ====================================================

# -*- coding:utf8 -*-
# Powered by Jinjing Li @2021-08-17 19:53:26


import os, sys, re
from ROOT import TFile, TH1D, TCanvas, TLegend

# topdir = "inputs/alphaN/rootfile/GeneratorInputs2025/DayaBay_PMT_pyrex/"
# topdir = "inputs/alphaN/rootfile/GeneratorInputs2025/DayaBay_GdLS/"
topdir = "inputs/alphaN/rootfile/GeneratorInputs2025/DayaBay_LS/"


def get_levels(nucleus, chain):
    res = []
    for rfile in os.listdir(topdir):
        if rfile.startswith(f"{nucleus}_{chain}"):
            state_str = rfile.split("_")[2]
            match = re.match(r"[a-zA-Z]+(\d+)", state_str)
            if match:
                number = int(match.group(1))  # 提取捕获组中的数字
                res.append(number)
    return sorted(list(set(res)))


source = sys.argv[1]  # U238/Th232
data = sys.argv[2]  # JENDL/TENDL
nucleus = sys.argv[3]  # na23/b11

states_list_U238 = {
    "b11": get_levels("b11", "U238"),
    "na23": get_levels("na23", "U238"),
    "b10": get_levels("b10", "U238"),
    "o18": get_levels("o18", "U238"),
    "si29": get_levels("si29", "U238"),
    "si30": get_levels("si30", "U238"),
    "c13": get_levels("c13", "U238"),
}

states_list_Th232 = {
    "b11": get_levels("b11", "Th232"),
    "na23": get_levels("na23", "Th232"),
    "b10": get_levels("b10", "Th232"),
    "o18": get_levels("o18", "Th232"),
    "si29": get_levels("si29", "Th232"),
    "si30": get_levels("si30", "Th232"),
    "c13": get_levels("c13", "Th232"),
}

states_list = (
    states_list_U238[nucleus] if source == "U238" else states_list_Th232[nucleus]
)

a = TCanvas("%s-%s" % (source, data), "%s-%s" % (source, data))
a.Draw()
a.SetRightMargin(0.03)
fnames = [
    topdir + "{3}_{0}Alpha_state{1}_{2}.root".format(source, idx, data, nucleus)
    for idx in (states_list)
]
afiles = [TFile(fname, "READ") for fname in fnames]
ahists = [afile["hspec"] for afile in afiles]
max_val = max([v.GetMaximum() for v in ahists])

ndraw = 0
for idx, ahist in enumerate(ahists):
    if ahist.Integral() > 0:
        ndraw += 1

low_legend_y = 0.85 - ndraw * 0.05
if low_legend_y < 0:
    low_legend_y = 0.1

lg0 = TLegend(0.8, low_legend_y, 0.95, 0.85)
lg0.SetTextSize(0.05)
lg0.SetTextFont(62)

for idx, ahist in enumerate(ahists):
    if ahist.Integral() > 0:
        ahist.Draw("" if idx == 0 else "same")
        lg0.AddEntry(ahist, "state-%d" % (states_list[idx]), "l")
    if idx == 0:
        ahist.SetMaximum(max_val * 1.2)
    ahist.SetXTitle("Neutron kinetic energy [MeV]")
    ahist.SetYTitle("Yield [/decay/0.01 MeV]")
    ahist.SetLineColor(idx + 1 if (idx + 1) % 5 != 0 else idx + 2)
    ahist.SetNdivisions()
    ahist.GetXaxis().CenterTitle()
    ahist.GetYaxis().CenterTitle()

if ndraw < 10:
    lg0.Draw("same")
# a.SetLogy()
a.Update()
a.Show()
a.SaveAs("./Plots/%s-%s-%s-generator-nspec.pdf" % (source, data, nucleus))
